# -*- coding:utf-8 -*-

from time import sleep

from pygene.gene import IntGene
from pygene.organism import Organism
from pygene.population import Population

from graph_presets import GraphPresets
from graph_utils import GraphWrapper

NODES_IN_GROUP = 3 # nodes in a group
PENALTY = 50000 # penalty for incorrect organisms
VARIANCE_EPS = 0.01

_graph = GraphWrapper(GraphPresets.nine_node())

class NodeGene(IntGene):
    """
    Gene which represents number of a node in a graph
    """
    randMin = 1    
    mutProb = 0.5
    mutAmt = 1


def tune_genetics(geneRandMax):
    """
    Tunes genetic parameters
    """
    NodeGene.randMax = geneRandMax
   
class Splitter(Organism):
    """
    Splitter organism
    """
    
    # preparing genome
    genome = {}
    for i in xrange(NODES_IN_GROUP):
        genome[str(i + 1)] = NodeGene
   
    def __init__(self, **kw):
        """
        Constructor
        """
        Organism.__init__(self, **kw)
        
        self._graph = _graph
                
    def fitness(self):
        """
        Fitness function. For correct gene values gives estimate value, calculated basing on a graph.
        For incorrect values, e.g. (1, 1, ...) gives penalty in a form of a huge value. I mean, they
        are born to be dead. hehe.
        """
        
        nodes = set()
        for i in xrange(NODES_IN_GROUP):
            nodes.add(_graph.node_map()[self[str(i + 1)]])
        
        if len(nodes) != NODES_IN_GROUP:
            return PENALTY
        
        estimate = self._estimate(self._graph, nodes)
        return float(estimate['outer'] - estimate['inner'])

    def _estimate(self, graph, nodes):
        """
        Calculates estimate for node group within a graph
        """
        assert isinstance(nodes, set)
        
        inner_visited = set()
        inner_weight = 0
        outer_weight = 0
        
        for start_node in nodes:
            neighbors = set(graph.graph.neighbors(start_node))
            inner = neighbors.intersection(nodes) - inner_visited 
            outer = neighbors - nodes
            
            for end_node in inner:
                inner_weight += graph.graph.edge_weight((start_node, end_node))
            inner_visited.add(start_node)
            
            for end_node in outer:
                outer_weight += graph.graph.edge_weight((start_node, end_node))
                
        return {"inner":inner_weight, "outer":outer_weight}

    
    def __repr__(self):
        """
        Defines how our Splitter should be dumped
        """
        coords = []
        for i in xrange(NODES_IN_GROUP):
            coords.append(str(self[str(i + 1)]))
        return "<Splitter fitness=%s (%s)>" % (self.fitness(),
                                               ", ".join(coords))
        
def populate_them():
    """
    Main routine
    """
    pop = make_population()
    while True:
        try:        
            pop.gen()
            print "START POPULATION"
            for partitioner in pop:
                print partitioner
            print "END POPULATION"
            sleep(1)
        except KeyboardInterrupt:
            pass
        
def mean_fitness(population):
    """
    Returns population's mean fitness
    """
    result = 0
    for organism in population:
        result += organism.fitness()
    result /= len(population)
    return result
    
def fitness_variance(population, mean=None):
    """
    Returns population's fitness variance
    """
    result = 0
    
    if mean is None:
        mean = mean_fitness(population)
    
    for partitioner in population:
        result += (partitioner.fitness() - mean) ** 2
    result /= len(population)
    return result

def make_population():
    """
    Makes population
    """
    total_nodes = _graph.node_count()
    return Population(species=Splitter, init=total_nodes,
                      childCount=total_nodes,
                      childCull=9, numNewOrganisms=5)
        
def populate_and_analyze():
    """
    Populates and analyzes resulting graph parts
    """
    
    tune_genetics(_graph.node_count())
    pop = make_population()
    
    while _graph.node_count() != 0:
        try:        
            pop.gen()
            
            mean = mean_fitness(pop) 
            variance = fitness_variance(pop, mean)
            
            print "Population mean and variance: ", mean, variance
            
            if variance <= VARIANCE_EPS:
                best_one = pop.best()
                
                group = []                
                for i in xrange(NODES_IN_GROUP):
                    group.append(_graph.node_map()[best_one[str(i + 1)]])
                    _graph.graph.del_node(_graph.node_map()[best_one[str(i + 1)]])
                    
                print "Group:", group
                
                _graph.reset_node_map()
                tune_genetics(_graph.node_count())
                pop = make_population()

            sleep(1)
        except KeyboardInterrupt:
            pass
    
        
def main():
    """
    Main routine
    """
    populate_and_analyze()
    
    
if __name__ == '__main__':
    main()